#encoding: utf8
from ctypes import *
from my_debugger_defines import *
import logging

kernel32 = windll.kernel32

fmt = '%(asctime)s %(levelname)s line:%(lineno)d]%(message)s'
formatter = logging.Formatter(fmt)
consoleHandler = logging.StreamHandler()
consoleHandler.setFormatter(formatter)
logging.getLogger('mydebugger').addHandler(consoleHandler)
logger = logging.getLogger('mydebugger')
logger.setLevel(logging.DEBUG)


class debugger():
    def __init__(self):
        self.h_process       = None
        self.pid             = None
        self.debugger_active = False 
        self.h_thread = None
        self.context = None
        self.exception = None
        self.exception_address = None
        self.breakpoints = {}
        self.first_breakpoints = True
        self.hardware_breakpoints ={}

        system_info = SYSTEM_INFO()
        kernel32.GetSystemInfo(byref(system_info))
        self.page_size = system_info.dwPageSize
        self.guarded_pages = {}
        self.memory_breakpoints = {}


    def load(self, path_to_exe):

        creation_flags = DEBUG_PROCESS

        startupinfo = STARTUPINFO()
        process_information = PROCESS_INFORMATION()

        startupinfo.dwFlags = 0x1
        startupinfo.wShowWindow = 0x0
        startupinfo.cb = sizeof(startupinfo)

        if kernel32.CreateProcessA(path_to_exe,
                    None,
                    None,
                    None,
                    None,
                    creation_flags,
                    None,
                    None,
                    byref(startupinfo),
                    byref(process_information)
                ):
            logger.info("[*] We have successfully launched the process!") 
            logger.info("[*] The Process ID I have is: %d" % \
                         process_information.dwProcessId)
        else:
            logger.info("[*] Error with error code %d." % kernel32.GetLastError())

    def open_process(self, pid):
        h_process = kernel32.OpenProcess(PROCESS_ALL_ACCESS, False, pid)
        return h_process

    

    def run(self):
        while self.debugger_active == True:
            self.get_debug_event()

    def get_debug_event(self):
        debug_event = DEBUG_EVENT()
        continue_status = DBG_CONTINUE
        # logger.debug("before wait for debug event")
        if kernel32.WaitForDebugEvent(byref(debug_event), 100):
            
            logger.info("thread id: %d" %(debug_event.dwThreadId))
            self.h_thread = self.open_thread(debug_event.dwThreadId)
            self.context = self.get_thread_context(h_thread=self.h_thread)
            # kernel32.CloseHandle(self.h_thread)

            # logger.debug("debug event code: %d, Thread ID: %d" % (
            #         debug_event.dwDebugEventCode,
            #         debug_event.dwThreadId)
            # )

            if debug_event.dwDebugEventCode == EXCEPTION_DEBUG_EVENT:
                self.exception = debug_event.u.Exception.ExceptionRecord.ExceptionCode
                self.exception_address = debug_event.u.Exception.ExceptionRecord.ExceptionAddress
                
                # call the internal handler for the exception event that just occured.
                if self.exception == EXCEPTION_ACCESS_VIOLATION:
                    logger.info("Access Violation Detected.") 
                elif self.exception == EXCEPTION_BREAKPOINT:
                    continue_status = self.exception_handler_breakpoint()
                elif self.exception == EXCEPTION_GUARD_PAGE:
                    logger.info("Guard Page Access Detected.") 
                elif self.exception == EXCEPTION_SINGLE_STEP:
                    continue_status = self.exception_handler_single_step()

            kernel32.ContinueDebugEvent(debug_event.dwProcessId,
                    debug_event.dwThreadId,
                    continue_status
            )
            # logger.debug("debug continue" )

    def exception_handler_breakpoint(self):
        logger.info("[*] Inside the breakpoint handler.")
        logger.info("[*] Exception address: 0x%08x" % self.exception_address) 
        return DBG_CONTINUE

    def exception_handler_single_step(self):
        logger.info("[*] Inside the single step handler.")
        logger.info("[*] Exception address: 0x%08x" % self.exception_address) 
        # Comment from PyDbg:
        # determine if this single step event occured in reaction to a hardware breakpoint and grab the hit breakpoint.
        # according to the Intel docs, we should be able to check for the BS flag in Dr6. but it appears that windows
        # isn't properly propogating that flag down to us.
        if self.context.Dr6 & 0x1 and self.hardware_breakpoints.has_key(0):
            slot = 0

        elif self.context.Dr6 & 0x2 and self.hardware_breakpoints.has_key(1):
            slot = 1
        elif self.context.Dr6 & 0x4 and self.hardware_breakpoints.has_key(2):
            slot = 2
        elif self.context.Dr6 & 0x8 and self.hardware_breakpoints.has_key(3):
            slot = 3
        else:
            # This wasn't an INT1 generated by a hw breakpoint
            continue_status = DBG_EXCEPTION_NOT_HANDLED

        # Now let's remove the breakpoint from the list
        if self.bp_del_hw(slot):
            continue_status = DBG_CONTINUE

        logger.info("[*] Hardware breakpoint removed.")
        return continue_status

    def attach(self, pid):
        self.h_process = self.open_process(pid)

        if kernel32.DebugActiveProcess(pid):
            self.debugger_active = True
            self.pid = int(pid)
            # self.run()
        else:
            logger.info("[*] Unable to attach to the process.")

    def detach(self):
        if kernel32.DebugActiveProcessStop(self.pid):
            logger.info("[*] Finished debugging. Exiting...")
            return True
        else:
            logger.info("There was an error")
            return False

    def open_thread(self, thread_id):
        h_thread = kernel32.OpenThread(THREAD_ALL_ACCESS, 
                    None, thread_id)

        if h_thread is not None:
            return h_thread
        else:
            logger.error("[*] Could not obtain a valid thread handle.")
            return False

    # 返回该进程的所有线程
    def enumerate_threads(self):
        thread_entry = THREADENTRY32()
        thread_list = []
        snapshot = kernel32.CreateToolhelp32Snapshot(
                    TH32CS_SNAPTHREAD, 
                    self.pid
                )
        if snapshot is not None:
            thread_entry.dwSize = sizeof(thread_entry)
            success = kernel32.Thread32First(
                        snapshot,
                        byref(thread_entry)
                    )
            while success:
                if thread_entry.th32OwnerProcessID == self.pid:
                    thread_list.append(thread_entry.th32ThreadID)
                success = kernel32.Thread32Next(
                        snapshot, 
                        byref(thread_entry)
                        )

            kernel32.CloseHandle(snapshot)
            return thread_list
        else:
            return False

    def get_thread_context(self, thread_id = None, h_thread = None):
        context = CONTEXT()
        context.ContextFlags = CONTEXT_FULL | CONTEXT_DEBUG_REGISTERS

        if h_thread is None:
            self.h_thread = self.open_thread(thread_id)

        if kernel32.GetThreadContext(self.h_thread, byref(context)):
            kernel32.CloseHandle(self.h_thread)
            return context
        else: 
            logger.error("[*] Get thread Context failed! Error code %d." % kernel32.GetLastError())
            
            return False

    def read_process_memory(self, address, length):
        data = ""
        read_buf = create_string_buffer(length)
        count = c_ulong(0)
        if not kernel32.ReadProcessMemory(self.h_process,
            address, read_buf, length, byref(count)):
            return False
        else:
            data += read_buf.raw
            return data


    def write_process_memory(self, address, data):
        count = c_ulong(0)
        length = len(data)

        c_data = c_char_p(data[count.value:])

        if not kernel32.WriteProcessMemory(self.h_process,
            address, c_data, length, byref(count)):
            return False
        else:
            return True

    def bp_set(self, address):
        if not self.breakpoints.has_key(address):
            try:
                original_byte = self.read_process_memory(address, 1)

                self.write_process_memory(address, "\xCC")

                self.breakpoints[address] = (address, original_byte)
                return True
            except:
                return False

    def func_resolve(self, dll, function):
        handler = kernel32.GetModuleHandleA(dll)
        address = kernel32.GetProcAddress(handler, function)
        kernel32.CloseHandle(handler)
        return address

    def bp_set_hw(self, address, length, condition):
        # 检测硬件断点的长度是否有效
        if length not in (1, 2, 4):
            return False
        else:
            length -= 1

        # 检测硬件断点的触发条件是否有效
        if condition not in (HW_ACCESS, HW_EXECUTE, HW_WRITE):
            return False
        # Check for available slots
        if not self.hardware_breakpoints.has_key(0):
            available = 0
        elif not self.hardware_breakpoints.has_key(1):
            available = 1
        elif not self.hardware_breakpoints.has_key(2):
            available = 2
        elif not self.hardware_breakpoints.has_key(3):
            available = 3
        else:
            return False

        # We want to set the debug register in every thread
        for thread_id in self.enumerate_threads():
            context = self.get_thread_context(thread_id=thread_id)
            if not context:
                continue
            # Enable the appropriate flag in the DR7
            # register to set the breakpoint
            context.Dr7 |= 1 << (available * 2)

            # Save the address of the breakpoint in the
            # free register that we found
            if   available == 0: context.Dr0 = address
            elif available == 1: context.Dr1 = address
            elif available == 2: context.Dr2 = address
            elif available == 3: context.Dr3 = address

            # Set the breakpoint condition
            context.Dr7 |= condition << ((available * 4) + 16)

            # Set the length
            context.Dr7 |= length << ((available * 4) + 18)

            # Set this threads context with the debug registers
            # set
            h_thread = self.open_thread(thread_id)
            kernel32.SetThreadContext(h_thread,byref(context))

        # update the internal hardware breakpoint array at the used slot index.
        self.hardware_breakpoints[available] = (address,length,condition)

        return True

    def bp_del_hw(self,slot):
        
        # Disable the breakpoint for all active threads
        for thread_id in self.enumerate_threads():
            
            context = self.get_thread_context(thread_id=thread_id)
            if not context:
                continue
            # Reset the flags to remove the breakpoint
            context.Dr7 &= ~(1 << (slot * 2))

            # Zero out the address
            if   slot == 0: 
                context.Dr0 = 0x00000000
            elif slot == 1: 
                context.Dr1 = 0x00000000
            elif slot == 2: 
                context.Dr2 = 0x00000000
            elif slot == 3: 
                context.Dr3 = 0x00000000

            # Remove the condition flag
            context.Dr7 &= ~(3 << ((slot * 4) + 16))

            # Remove the length flag
            context.Dr7 &= ~(3 << ((slot * 4) + 18))

            # Reset the thread's context with the breakpoint removed
            h_thread = self.open_thread(thread_id)
            kernel32.SetThreadContext(h_thread,byref(context))
            
        # remove the breakpoint from the internal list.
        del self.hardware_breakpoints[slot]

        return True

    def bp_set_mem(self, address, size):
        mbi = MEMORY_BASIC_INFORMATION()

        if kernel32.VirtualQueryEx(self.h_process, 
                address, byref(mbi), sizeof(mbi)) < sizeof(mbi):
            return False

        current_page = mbi.BaseAddress
        while current_page <= address + size:
            self.guarded_pages.append(current_page)
            old_protection = c_ulong(0)
            if not kernel32.VirtualProtectEx(self.h_process,
                    current_page, size, 
                    mbi.Protect | PAGE_GUARD, byref(old_protection)):
                return False
            current_page += self.page_size

        self.memory_breakpoints[address] = (address, size, mbi)
        return True

    